"""
实现数据在传输的过程中，充分利用带宽，且不造成网络拥堵的软件程序，麻烦用java实现
"""

"""
CongestionControl 类：实现了 TCP 拥塞控制算法，包括慢启动和拥塞避免机制
  · 根据 ACK 确认和超时事件动态调整拥塞窗口大小
  · 记录拥塞窗口变化历史用于可视化
  · Packet 类：表示网络数据包，包含序列号、数据内容和时间戳
NetworkSimulator 类：模拟网络环境
  · 支持设置带宽、延迟和丢包率
  · 模拟数据包传输和随机丢包
  · 计算并记录带宽利用率
BandwidthController 类：智能带宽控制器
  · 根据拥塞窗口动态调整发送速率
  · 实现数据包的发送、确认处理和超时重传
  · 记录并可视化带宽使用情况

使用示例展示了如何模拟数据传输过程，并在中途改变网络条件（如降低带宽、增加丢包率），
观察系统如何自适应调整发送速率以充分利用可用带宽同时避免网络拥塞。程序会生成两个图表：
带宽利用率随时间的变化和拥塞窗口大小的动态调整过程。
这个实现基于 TCP 拥塞控制的原理，但做了简化处理，适用于教学和模拟目的。在实际应用中，可能需要更复杂的算法和更精确的网络参数调整。
"""

import time
import threading
import random
import queue
from collections import deque, defaultdict
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import numpy as np


class CongestionControl:
    """拥塞控制算法实现"""

    def __init__(self):
        # 初始窗口大小 (MSS)
        self.cwnd = 1.0
        # 慢启动阈值
        self.ssthresh = 65535
        # 上次确认的序列号
        self.last_ack = 0
        # 拥塞窗口历史记录（用于可视化）
        self.cwnd_history = []
        # 时间戳历史记录（用于可视化）
        self.time_history = []
        # 开始时间
        self.start_time = time.time()

    def on_ack_received(self, ack_num):
        """处理确认包"""
        current_time = time.time() - self.start_time

        # 记录历史数据
        self.cwnd_history.append(self.cwnd)
        self.time_history.append(current_time)

        # 检测是否有新确认
        if ack_num > self.last_ack:
            self.last_ack = ack_num

            # 慢启动和拥塞避免
            if self.cwnd < self.ssthresh:
                # 慢启动：每个ACK增加1个MSS
                self.cwnd += 1.0
            else:
                # 拥塞避免：每个RTT增加1个MSS
                self.cwnd += 1.0 / self.cwnd

    def on_timeout(self):
        """处理超时事件"""
        current_time = time.time() - self.start_time

        # 记录历史数据
        self.cwnd_history.append(self.cwnd)
        self.time_history.append(current_time)

        # 调整拥塞窗口和阈值
        self.ssthresh = max(2.0, self.cwnd / 2)
        self.cwnd = 1.0

    def get_cwnd(self):
        """获取当前拥塞窗口大小"""
        return self.cwnd


class Packet:
    """表示网络数据包"""

    def __init__(self, seq_num, data, timestamp=None):
        self.seq_num = seq_num  # 序列号
        self.data = data  # 数据包内容
        self.timestamp = timestamp or time.time()  # 发送时间戳
        self.retransmitted = False  # 是否为重传包


class NetworkSimulator:
    """网络模拟器，模拟数据包传输和网络拥塞"""

    def __init__(self, bandwidth_mbps=10, delay_ms=50, loss_rate=0.01):
        self.bandwidth_mbps = bandwidth_mbps  # 带宽 (Mbps)
        self.delay_ms = delay_ms  # 延迟 (ms)
        self.loss_rate = loss_rate  # 丢包率
        self.queue = queue.Queue()  # 数据包队列
        self.is_running = False
        self.thread = None
        self.receive_callback = None
        self.transmission_history = defaultdict(list)  # 传输历史记录

    def start(self):
        """启动网络模拟器"""
        if not self.is_running:
            self.is_running = True
            self.thread = threading.Thread(target=self._transmit_packets)
            self.thread.daemon = True
            self.thread.start()

    def stop(self):
        """停止网络模拟器"""
        self.is_running = False
        if self.thread:
            self.thread.join(timeout=1.0)

    def send_packet(self, packet):
        """发送数据包"""
        # 计算传输延迟 (基于带宽)
        data_size_bytes = len(packet.data)
        transmission_delay = (data_size_bytes * 8) / (self.bandwidth_mbps * 1e6)  # 秒

        # 记录传输信息
        send_time = time.time()
        self.transmission_history[send_time].append((packet.seq_num, data_size_bytes))

        # 将数据包放入队列，添加传输延迟
        delivery_time = send_time + transmission_delay + (self.delay_ms / 1000)
        self.queue.put((delivery_time, packet))

    def set_receive_callback(self, callback):
        """设置接收回调函数"""
        self.receive_callback = callback

    def _transmit_packets(self):
        """数据包传输线程"""
        while self.is_running:
            try:
                # 获取队列中的下一个数据包
                delivery_time, packet = self.queue.get(timeout=0.1)

                # 检查是否应该发送这个数据包
                if time.time() >= delivery_time:
                    # 模拟丢包
                    if random.random() >= self.loss_rate:
                        # 调用接收回调
                        if self.receive_callback:
                            self.receive_callback(packet)
                    self.queue.task_done()
                else:
                    # 还没到发送时间，放回队列
                    self.queue.put((delivery_time, packet))

            except queue.Empty:
                continue

    def get_bandwidth_utilization(self, window_size=1.0):
        """计算带宽利用率 (window_size秒窗口)"""
        now = time.time()
        total_bytes = 0

        for timestamp, packets in self.transmission_history.items():
            if timestamp >= now - window_size:
                for _, size in packets:
                    total_bytes += size

        # 转换为Mbps
        return (total_bytes * 8) / (window_size * 1e6)


class BandwidthController:
    """智能带宽控制器"""

    def __init__(self, network_simulator, packet_size=1024):
        self.network = network_simulator
        self.packet_size = packet_size  # 数据包大小 (bytes)
        self.congestion_control = CongestionControl()
        self.send_queue = queue.Queue()
        self.sent_packets = {}  # 已发送但未确认的数据包
        self.next_seq_num = 0
        self.is_running = False
        self.send_thread = None
        self.receive_thread = None
        self.ack_queue = queue.Queue()
        self.bandwidth_history = []
        self.time_history = []
        self.start_time = time.time()

        # 设置网络接收回调
        self.network.set_receive_callback(self._receive_packet)

    def start(self):
        """启动带宽控制器"""
        if not self.is_running:
            self.is_running = True
            self.start_time = time.time()
            self.send_thread = threading.Thread(target=self._sender_loop)
            self.send_thread.daemon = True
            self.send_thread.start()

            self.receive_thread = threading.Thread(target=self._receiver_loop)
            self.receive_thread.daemon = True
            self.receive_thread.start()

    def stop(self):
        """停止带宽控制器"""
        self.is_running = False
        if self.send_thread:
            self.send_thread.join(timeout=1.0)
        if self.receive_thread:
            self.receive_thread.join(timeout=1.0)

    def enqueue_data(self, data):
        """将数据加入发送队列"""
        # 分割数据为固定大小的数据包
        for i in range(0, len(data), self.packet_size):
            chunk = data[i:i + self.packet_size]
            self.send_queue.put(chunk)

    def _sender_loop(self):
        """发送循环，根据拥塞窗口控制发送速率"""
        while self.is_running:
            try:
                # 检查拥塞窗口大小
                cwnd = self.congestion_control.get_cwnd()

                # 计算当前可以发送的数据包数量
                packets_in_flight = len(self.sent_packets)
                available_slots = max(0, int(cwnd) - packets_in_flight)

                # 发送数据包
                for _ in range(available_slots):
                    try:
                        data = self.send_queue.get_nowait()
                        seq_num = self.next_seq_num
                        self.next_seq_num += 1

                        packet = Packet(seq_num, data)
                        self.sent_packets[seq_num] = packet
                        self.network.send_packet(packet)

                        # 记录带宽使用情况
                        current_time = time.time() - self.start_time
                        self.bandwidth_history.append(self.network.get_bandwidth_utilization())
                        self.time_history.append(current_time)

                    except queue.Empty:
                        # 没有更多数据可发送
                        break

                # 检查超时
                current_time = time.time()
                for seq_num, packet in list(self.sent_packets.items()):
                    if current_time - packet.timestamp > 1.0:  # 假设超时时间为1秒
                        # 重传超时的数据包
                        packet.retransmitted = True
                        packet.timestamp = current_time
                        self.network.send_packet(packet)
                        # 通知拥塞控制发生超时
                        self.congestion_control.on_timeout()

                # 控制发送频率
                time.sleep(0.01)

            except Exception as e:
                print(f"发送循环异常: {e}")
                time.sleep(0.1)

    def _receive_packet(self, packet):
        """接收数据包的回调函数"""
        # 检查是否为ACK包 (这里简化处理，假设所有包都是ACK)
        if packet.seq_num in self.sent_packets:
            self.ack_queue.put(packet.seq_num)

    def _receiver_loop(self):
        """接收确认包的循环"""
        while self.is_running:
            try:
                # 获取下一个确认号
                ack_num = self.ack_queue.get(timeout=0.1)

                # 处理确认
                if ack_num in self.sent_packets:
                    del self.sent_packets[ack_num]
                    # 通知拥塞控制收到确认
                    self.congestion_control.on_ack_received(ack_num)

            except queue.Empty:
                continue
            except Exception as e:
                print(f"接收循环异常: {e}")
                time.sleep(0.1)

    def plot_bandwidth_usage(self):
        """绘制带宽使用情况图表"""
        plt.figure(figsize=(12, 6))

        # 绘制带宽使用情况
        plt.subplot(2, 1, 1)
        plt.plot(self.time_history, self.bandwidth_history)
        plt.axhline(y=self.network.bandwidth_mbps, color='r', linestyle='--', label='可用带宽')
        plt.xlabel('时间 (秒)')
        plt.ylabel('带宽使用 (Mbps)')
        plt.title('带宽利用率')
        plt.legend()
        plt.grid(True)

        # 绘制拥塞窗口大小
        plt.subplot(2, 1, 2)
        plt.plot(self.congestion_control.time_history, self.congestion_control.cwnd_history)
        plt.xlabel('时间 (秒)')
        plt.ylabel('拥塞窗口大小 (MSS)')
        plt.title('拥塞窗口动态变化')
        plt.grid(True)

        plt.tight_layout()
        plt.show()


# 使用示例
def simulate_network_transfer():
    # 创建网络模拟器 (10Mbps带宽, 50ms延迟, 1%丢包率)
    network = NetworkSimulator(bandwidth_mbps=10, delay_ms=50, loss_rate=0.01)

    # 创建带宽控制器
    controller = BandwidthController(network)

    try:
        # 启动网络和控制器
        network.start()
        controller.start()

        # 生成大量数据进行传输
        print("开始模拟数据传输...")
        data = b"X" * 1024 * 1024  # 1MB数据
        controller.enqueue_data(data)

        # 模拟网络条件变化
        print("模拟网络拥塞...")
        time.sleep(5)
        network.bandwidth_mbps = 5  # 降低带宽到5Mbps
        network.loss_rate = 0.05  # 增加丢包率到5%

        time.sleep(10)

        print("恢复网络条件...")
        network.bandwidth_mbps = 10  # 恢复带宽
        network.loss_rate = 0.01  # 恢复丢包率

        # 等待传输完成
        time.sleep(10)

    finally:
        # 停止控制器和网络
        controller.stop()
        network.stop()

        # 绘制带宽使用情况
        controller.plot_bandwidth_usage()


if __name__ == "__main__":
    simulate_network_transfer()



